import torch  # 导入PyTorch库
from MyData import MyDataset  # 导入自定义数据集类
from torch.utils.data import DataLoader  # 导入PyTorch的数据加载器
from net import Model  # 导入自定义的模型类
from transformers import BertTokenizer  # 导入AdamW优化器和Bert分词器
# Transformers 4.50.3 已移除 optimization.AdamW 的显式导出，或需依赖 PyTorch 的 torch.optim.AdamW 实现‌
from torch.optim import AdamW

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 设置设备为CUDA（如果可用），否则为CPU
EPOCH = 30000  # 设置训练的总轮数

token = BertTokenizer.from_pretrained(r"../../my_model_cache/bert-base-chinese/models--bert-base-chinese/snapshots/c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f")  # 加载预训练的Bert分词器

def collate_fn(data):
    sentes = [i[0] for i in data]  # 提取数据集中的句子
    label = [i[1] for i in data]  # 提取数据集中的标签
    # print(sentes)
    # 编码句子
    data = token.batch_encode_plus(batch_text_or_text_pairs=sentes,
                                   truncation=True,
                                   padding="max_length",
                                   max_length=1500,
                                   return_tensors="pt",
                                   return_length=True)  # 批量编码句子
    input_ids = data["input_ids"]  # 获取输入ID
    attention_mask = data["attention_mask"]  # 获取注意力掩码
    token_type_ids = data["token_type_ids"]  # 获取token类型ID
    labels = torch.LongTensor(label)  # 将标签转换为LongTensor
    # print(input_ids,attention_mask,token_type_ids)
    return input_ids, attention_mask, token_type_ids, labels  # 返回处理后的数据

# 创建数据集
train_dataset = MyDataset("train")  # 创建训练数据集
val_dataset = MyDataset("validation")  # 创建验证数据集
# 创建数据加载器
train_laoder = DataLoader(dataset=train_dataset,
                          batch_size=20,
                          shuffle=True,
                          drop_last=True,
                          collate_fn=collate_fn)  # 为训练数据集创建数据加载器

if __name__ == '__main__':
    # 开始训练
    print(DEVICE)  # 打印使用的设备
    model = Model().to(DEVICE)  # 初始化模型并将其移动到指定设备
    optimizer = AdamW(model.parameters(), lr=5e-4)  # 使用AdamW优化器
    loss_func = torch.nn.CrossEntropyLoss()  # 使用交叉熵损失函数

    model.train()  # 设置模型为训练模式
    for epoch in range(EPOCH):  # 训练EPOCH轮
        sum_val_acc = 0
        sum_val_loss = 0
        # 训练
        for i, (input_ids, attention_mask, token_type_ids, labels) in enumerate(train_laoder):
            # print(input_ids)
            input_ids, attention_mask, token_type_ids, labels = input_ids.to(DEVICE), attention_mask.to(
                DEVICE), token_type_ids.to(DEVICE), labels.to(DEVICE)  # 将数据移动到指定设备
            out = model(input_ids, attention_mask, token_type_ids)  # 前向传播

            loss = loss_func(out, labels)  # 计算损失

            optimizer.zero_grad()  # 清空优化器的梯度
            loss.backward()  # 反向传播
            optimizer.step()  # 更新模型参数

            if i % 5 == 0:  # 每5个批次打印一次信息
                out = out.argmax(dim=1)  # 获取预测结果
                acc = (out == labels).sum().item() / len(labels)  # 计算准确率
                print(epoch, i, loss.item(), acc)  # 打印轮数、批次、损失和准确率

        torch.save(model.state_dict(), f"params/{epoch}bert-weibo.pth")  # 保存模型参数
        print(epoch, "参数保存成功！")  # 打印成功信息
